package com.feidee.fd.sml.algorithm.component.ml.classification

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.LogisticRegression

/**
  * @Author: dongguosheng, songhaicheng
  * @Date: 2018/8/16 16:03
  */
case class LogisticRegressionParam(
                                    override val input_pt: String,
                                    override val output_pt: String,
                                    override val hive_table: String,
                                    override val flow_time: String,
                                    override val modelPath: String,
                                    override var labelCol: String,
                                    override val featuresCol: String,
                                    override var predictionCol: String,
                                    override val metrics: Array[String],
                                    // 弹性网络混合参数。[0.0, 1.0]，默认 0.0（对于alpha = 0，惩罚是 L2 惩罚。 对于 alpha = 1，它是 L1 惩罚。 对于（0,1）中的 α，惩罚是 L1 和 L2 的组合）
                                    elasticNet: Double,
                                    // 正则化系数，>= 0，默认 0.0
                                    regParam: Double,
                                    // Spark treeAggregate 算子的参数，适度增大可以加快训练速度，>= 2，默认 2
                                    aggregationDepth: Int,
                                    // 标签列的值分布类型，支持 [auto, binomial, multinomial], 默认 auto（自动根据标签数选择）
                                    family: String,
                                    // 是否需要计算截距，默认 true
                                    fitIntercept: Boolean,
                                    // 是否对特征进行标准化操作，默认 true
                                    standardization: Boolean,
                                    // 最大迭代次数，默认 100
                                    maxIter: Int,
                                    // 优化算法迭代求解过程的收敛阀值，>= 0，默认 1E-6
                                    tol: Double,
                                    // 样本权重列
                                    weightCol: String,
                                    // 预测概率列名，默认 probability
                                    probabilityCol: String,
                                    // 原始结果列，默认 rawPrediction
                                    rawPredictionCol: String,
                                    // 分类类别阈值，长度要与标签类别数一致
                                    thresholds: Array[Double],
                                    // 二分法分类阈值，[0, 1]，默认 0.5
                                    threshold: Double
                                  ) extends MLParam {
  def this() = this(null, null, null, null, null, "label", "features", "prediction", Array.empty[String], 0.0, 0.0, 2,
    "auto", true, true, 100, 1E-6, null, "probability", "rawPrediction", Array.empty[Double], 0.5)

  override def verify(): Unit = {
    super.verify()
    val families = Array("auto", "binomial", "multinomial")

    require(families.contains(family.toLowerCase),
      s"param family only accepts ${families.mkString("[", ", ", "]")}, but has $family")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("elasticNet" -> elasticNet)
    map += ("regParam" -> regParam)
    map += ("aggregationDepth" -> aggregationDepth)
    map += ("family" -> family)
    map += ("fitIntercept" -> fitIntercept)
    map += ("standardization" -> standardization)
    map += ("maxIter" -> maxIter)
    map += ("tol" -> tol)
    map += ("weightCol" -> weightCol)
    map += ("probabilityCol" -> probabilityCol)
    map += ("rawPredictionCol" -> rawPredictionCol)
    map += ("thresholds" -> thresholds)
    map += ("threshold" -> threshold)
    map
  }
}


class LRComponent extends AbstractClassificationComponent[LogisticRegressionParam] {

  override def setUp(param: LogisticRegressionParam): PipelineStage = {
    val lr = new LogisticRegression()
      .setLabelCol(param.labelCol)
      .setFeaturesCol(param.featuresCol)
      .setRawPredictionCol(param.rawPredictionCol)
      .setPredictionCol(param.predictionCol)
      .setProbabilityCol(param.probabilityCol)
      .setElasticNetParam(param.elasticNet)
      .setRegParam(param.regParam)
      .setAggregationDepth(param.aggregationDepth)
      .setFitIntercept(param.fitIntercept)
      .setStandardization(param.standardization)
      .setMaxIter(param.maxIter)
      .setFamily(param.family)
      .setThreshold(param.threshold)
      .setTol(param.tol)

    if (tool.isNotNull(param.weightCol)) {
      lr.setWeightCol(param.weightCol)
    }
    if (param.thresholds.nonEmpty) {
      lr.setThresholds(param.thresholds)
    }

    lr
  }

}

object LRComponent {
  def apply(paramStr: String): Unit = {
    new LRComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    LRComponent(args(0))
  }

}
